import argparse
import os
import cv2
from medpy.io import load,save
import numpy as np
def make_parser():
    parser = argparse.ArgumentParser("nii dataset parser")
    # distributed
    parser.add_argument(
        "--datapath", default="", type=str, help="data_nii_path"
    )
    parser.add_argument(
        "--savepath", default="", type=str, help="save path"
    )
    return parser
    
    
    
def label_seg(nii_path,save_folder):
    data_nii_path = nii_path + '/data'
    label_nii_path = nii_path+'/label'
    file_list = os.listdir(label_nii_path)    #批量处理
    for nii_label in file_list:

        liver_tumor_label ,header= load(os.path.join(label_nii_path, nii_label))   # 加载nii文件
        liver_tumor_data ,header1= load(os.path.join(data_nii_path, nii_label.replace('segmentation', 'volume')))   # 加载nii文件
        liver_tumor_data[liver_tumor_data > 250] = 250	#加窗处理，过滤掉过于的组织
        liver_tumor_data[liver_tumor_data < -250] = -250
        liver = np.zeros(liver_tumor_label.shape)  # 创建一个全为0的矩阵，形状与加载的nii文件大小一样
        tumor = np.zeros(liver_tumor_label.shape)
        index = np.where(liver_tumor_label == 1)   # 肝脏标签为1，返回索引位置
        liver[index] = 1    # 将得到的肝脏坐标在新的矩阵中置为1，得到只包含肝脏的标签
        index1 = np.where(liver_tumor_label == 2)  # 肿瘤标签为2，返回索引位置
        tumor[index1] = 1   # 将得到的肿瘤坐标在新的矩阵中置为1，得到只包含肿瘤的标签
        slice_num = liver_tumor_label.shape[2]     # 获取切片数量
        tumor_path = save_folder+'/tumor_label'     # 分离的标签设置保存路径
        liver_path = save_folder+'/liver_label'
        train_tumor_path = save_folder+'/train_tumor'     # 训练数据保存路径
        train_liver_path = save_folder+'/train_liver'
        if not os.path.exists(tumor_path):      # 判断文件目录是否存在，不存在则创建
            os.mkdir(tumor_path)
        if not os.path.exists(liver_path):
            os.mkdir(liver_path)
        if not os.path.exists(train_tumor_path):      # 判断文件目录是否存在，不存在则创建
            os.mkdir(train_tumor_path)
        if not os.path.exists(train_liver_path):
            os.mkdir(train_liver_path)
        
        img_name = os.path.split(nii_label)     # 分离文件名，用以保存时的文件名前缀
        img_name = img_name[-1]
        img_name = img_name.split('.')
        img_name = img_name[0]
        for i in range(0, slice_num):

            if liver[:, :, i].max() > 0:    # 肝脏标签图片
                label_liver = liver[:, :, i] * 122
                liver_train_data = (liver_tumor_data[:, :, i] - liver_tumor_data[:, :, i].min()) / (liver_tumor_data[:, :, i].max() - liver_tumor_data[:, :, i].min()) * 255
                print(cv2.imwrite("%s/%s-%d.png" % (liver_path, "liver-"+img_name, i), label_liver))	
                print(cv2.imwrite("%s/%s-%d.png" % (train_liver_path, img_name.replace('segmentation', 'volume'), i), liver_train_data))	#保持分割出来的标签与数据的序号对应，因为肿瘤的标签量比较少
            if tumor[:, :, i].max() > 0:    # 肿瘤标签图片
                label_tumor = tumor[:, :, i] * 122
                tumor_train_data = (liver_tumor_data[:, :, i] - liver_tumor_data[:, :, i].min()) / (liver_tumor_data[:, :, i].max() - liver_tumor_data[:, :, i].min()) * 255
                print(cv2.imwrite("%s/%s-%d.png" % (tumor_path, "tumor-"+img_name, i), label_tumor))
                print(cv2.imwrite("%s/%s-%d.png" % (train_tumor_path, img_name.replace('segmentation', 'volume'), i),
                                  tumor_train_data))	#保持分割出来的标签与数据的序号对应，因为肿瘤的标签量比较少
    # save(liver, os.path.join(save_folder,"liver.nii"))
    # save(tumor, os.path.join(save_folder, "tumor.nii"))

if __name__ == "__main__":
    args = make_parser().parse_args()
    label_nii_path = args.datapath

    save_path = args.savepath

    label_seg(label_nii_path,save_path)
